#!/bin/bash

# Copyright (c) 2024, Kuaishou Technology. All rights reserved.

set -euo pipefail

source ./llama-7b

export SEQ_LENGTH=12288

export HOSTFILE=$1
NUM_NODE=$2
export MASTER_ADDR=$(cat $HOSTFILE | head -n 1 | sed -s 's/slots=8//g')
export NUM_GPUS=$((NUM_NODE*8))

export GLOBAL_BATCH_SIZE=$((NUM_NODE*8*4))
export TP=1
export CP=1
export PP=4
export PP_l=1
export CKPT=ours
export OFFLOAD_ALPHA=0.40

./pretrain_llama.sh
